import argparse
from typing import List

import gradio as gr

import iflytech_assistant.assistant.heqc as heqc
import iflytech_assistant.assistant.rag_heqc as rag_heqc

parser = argparse.ArgumentParser()
parser.add_argument("-tf", "--tags-file", type=str, help="tags file path, xlsx format")
args = parser.parse_args()

tag_objs, refine_words = heqc.load_tags_from_excel(args.tags_file)


def on_target_change(target: str):
    # change visible of tag
    return gr.update(choices=[tag.tag for tag in tag_objs[target]]), gr.update(
        choices=refine_words[target]
    )


def on_generate_click(user_input: str, target: str, mode: str, tags: List[int]):
    mode = "reply" if mode == "回复" else "polish"
    tags = [tag_objs[target][tag] for tag in tags]
    system_prompt, user_prompt, llm_response, suggestions = heqc.generate(
        user_input, target, tags, mode
    )
    state = {
        "target": target,
        "mode": mode,
        "user_input": user_input,
        "system_prompt": system_prompt,
        "user_prompt": user_prompt,
        "llm_response": llm_response,
        "suggestions": suggestions,
    }

    prompt_preview = (
        f"# SYSTEM_PROMPT：\n{system_prompt}\n\n# USER_PROMPT：\n{user_prompt}"
    )

    system_prompt, user_prompt, llm_response, rag_suggestions = rag_heqc.generate(
        user_input, target, tags, mode
    )
    rag_state = {
        "target": target,
        "mode": mode,
        "user_input": user_input,
        "system_prompt": system_prompt,
        "user_prompt": user_prompt,
        "llm_response": llm_response,
        "suggestions": rag_suggestions,
    }
    rag_prompt_preview = (
        f"# SYSTEM_PROMPT：\n{system_prompt}\n\n# USER_PROMPT：\n{user_prompt}"
    )
    return (
        str(suggestions),
        suggestions.explanations,
        prompt_preview,
        state,
        str(rag_suggestions),
        rag_suggestions.explanations,
        rag_prompt_preview,
        rag_state,
    )


def on_refine_click(state: dict, rag_state: dict, refine_word_index: int):
    if not state:
        gr.Warning("请点击生成，再进行修改。", duration=5)
        return "", "", "", {}

    target = state["target"]
    refine_suggestions, refine_prompt = heqc.refine(
        state["user_input"],
        state["system_prompt"],
        refine_words[target][refine_word_index],
        state["llm_response"],
    )
    state["llm_response"] = refine_prompt
    state["suggestions"] = refine_suggestions

    system_prompt = state["system_prompt"]
    user_prompt = state["user_prompt"]
    prompt_preview = f"# SYSTEM_PROMPT：\n{system_prompt}\n\n# USER_PROMPT：\n{user_prompt}\n\n# REFINED_PROMPT：\n{refine_prompt}"

    rag_refine_suggestions, rag_refine_prompt = rag_heqc.refine(
        rag_state["user_input"],
        rag_state["system_prompt"],
        refine_words[target][refine_word_index],
        rag_state["llm_response"],
    )
    rag_state["llm_response"] = rag_refine_prompt
    rag_state["suggestions"] = rag_refine_suggestions

    rag_system_prompt = rag_state["system_prompt"]
    rag_user_prompt = rag_state["user_prompt"]
    rag_prompt_preview = f"# SYSTEM_PROMPT：\n{rag_system_prompt}\n\n# USER_PROMPT：\n{rag_user_prompt}\n\n# REFINED_PROMPT：\n{rag_refine_prompt}"

    return (
        str(refine_suggestions),
        refine_suggestions.explanations,
        prompt_preview,
        state,
        str(rag_refine_suggestions),
        rag_refine_suggestions.explanations,
        rag_prompt_preview,
        rag_state,
    )


# Gradio 界面
with gr.Blocks() as demo:
    agent_state = gr.State(value={})
    rag_agent_state = gr.State(value={})

    gr.Markdown("# 高情商沟通RAG Demo")
    with gr.Row():
        with gr.Column(scale=1):
            suggestions = gr.Textbox(
                label="高情商建议", lines=6, interactive=False, scale=1
            )
            explainations = gr.Textbox(
                label="高情商解释", lines=6, interactive=False, scale=2
            )
            prompt_preview_button = gr.Button("查看提示词")
            prompt_preview = gr.Markdown(visible=False)

        with gr.Column(scale=1):
            rag_suggestions = gr.Textbox(
                label="高情商建议+RAG", lines=6, interactive=False, scale=1
            )
            rag_explainations = gr.Textbox(
                label="高情商解释+RAG", lines=6, interactive=False, scale=2
            )
            rag_prompt_preview_button = gr.Button("查看提示词")
            rag_prompt_preview = gr.Markdown(visible=False)

        with gr.Column(scale=1):
            user_input = gr.Textbox(label="用户输入", lines=1)
            with gr.Row():
                target = gr.Dropdown(
                    label="选择聊天对象", choices=list(tag_objs.keys())
                )
                mode = gr.Dropdown(
                    label="选择模式",
                    choices=["回复", "润色"],
                )

            tag = gr.Dropdown(
                label="选择偏好标签",
                choices=[],
                multiselect=True,
                value=None,
                interactive=True,
                max_choices=5,
                type="index",
            )
            generate_button = gr.Button("生成")
            refine_word = gr.Dropdown(
                choices=[],
                label="选择修改提示词",
                value=None,
                interactive=True,
                type="index",
            )
            refine_button = gr.Button("修改")

    target.change(fn=on_target_change, inputs=target, outputs=[tag, refine_word])
    generate_button.click(
        on_generate_click,
        inputs=[user_input, target, mode, tag],
        outputs=[
            # heqc
            suggestions,
            explainations,
            prompt_preview,
            agent_state,
            # rag_heqc
            rag_suggestions,
            rag_explainations,
            rag_prompt_preview,
            agent_state,
        ],
    )
    refine_button.click(
        on_refine_click,
        inputs=[agent_state, rag_agent_state, refine_word],
        outputs=[
            suggestions,
            explainations,
            prompt_preview,
            agent_state,
            rag_suggestions,
            rag_explainations,
            rag_prompt_preview,
            rag_agent_state,
        ],
    )
    prompt_preview_button.click(
        lambda: gr.update(visible=not prompt_preview.visible),
        outputs=[prompt_preview],
    )
    rag_prompt_preview_button.click(
        lambda: gr.update(visible=not rag_prompt_preview.visible),
        outputs=[rag_prompt_preview],
    )

demo.launch(share=True)
